#include "BasicRemoteInject.h"
#include <string>
#include <Shlwapi.h>
#include <Windows.h>
#pragma comment(lib,"Shlwapi.lib")


namespace utility
{
	BasicRemoteInject::BasicRemoteInject(void)
	{
	}

	BasicRemoteInject::~BasicRemoteInject(void)
	{
	}

	bool BasicRemoteInject::RemoteInjectLibrary( const unsigned long dwProcessId, const TCHAR* szLibrary )
	{
		if( dwProcessId==0 || !szLibrary )
		{
			return false;
		}

		if( !::PathFileExists(szLibrary) )
		{
			return false;
		}

		//DWORD pId;
		//HWND hWinPro=::FindWindow("ProgMan",NULL); 
		//::GetWindowThreadProcessId(hWinPro,&pId);

		DWORD THREADSIZE = 1024*2;//const DWORD THREADSIZE = 1024*4;
		HANDLE pRemoteThread = NULL;
		HANDLE hRemoteProcess = NULL;
		PTHREAD_START_ROUTINE pfnAddr = NULL;
		void *pFileRemote = NULL;

		hRemoteProcess=::OpenProcess(PROCESS_ALL_ACCESS,false,dwProcessId); 
		int nNum = 0;
		do 
		{
			THREADSIZE = THREADSIZE << 1;
			pFileRemote=::VirtualAllocEx(hRemoteProcess,0,THREADSIZE,MEM_COMMIT|MEM_RESERVE,PAGE_EXECUTE_READWRITE);
		} while (++nNum < 10 && !pFileRemote);

		if(pFileRemote && !::WriteProcessMemory(hRemoteProcess,pFileRemote,szLibrary/*"d:\\RemoteDll.dll"*/,THREADSIZE,NULL))
		{
			::VirtualFreeEx(hRemoteProcess, pFileRemote, THREADSIZE, MEM_RELEASE );
			return false;
		}

#ifdef UNICODE
		std::string strLoadLibrary = "LoadLibraryW";
#else
		std::string strLoadLibrary = "LoadLibraryA";
#endif
		pfnAddr=(PTHREAD_START_ROUTINE)GetProcAddress(GetModuleHandle(TEXT("Kernel32")),strLoadLibrary.c_str());

		pRemoteThread=::CreateRemoteThread(hRemoteProcess,NULL,0,pfnAddr,pFileRemote,0,NULL);

		if(pRemoteThread==NULL)
		{
			::VirtualFreeEx(hRemoteProcess, pFileRemote, THREADSIZE, MEM_RELEASE );
			return false;
		}

		::WaitForSingleObject(pRemoteThread,INFINITE);
		DWORD hLibModule;
		::GetExitCodeThread( pRemoteThread, &hLibModule);
		::CloseHandle( pRemoteThread );
		::VirtualFreeEx( hRemoteProcess, pFileRemote,THREADSIZE, MEM_RELEASE);

		SREMOTEINFO* pRemoteInfo = new SREMOTEINFO();
		pRemoteInfo->dwProcessId = dwProcessId;
		pRemoteInfo->hRemoteThread = pRemoteThread;
		pRemoteInfo->pVirtualAddress = pFileRemote;
		pRemoteInfo->hRemoteProcess = hRemoteProcess;
		pRemoteInfo->dwMemSize = THREADSIZE;
		pRemoteInfo->dwType = BasicRemoteInject::TYPE_DLL;
		pRemoteInfo->hLibModelAddress = hLibModule;
		_tcsncpy_s(pRemoteInfo->szFileName, MAX_PATH*3, szLibrary, _tcslen(szLibrary));
		this->m_mRemoteList.insert(std::make_pair(dwProcessId, pRemoteInfo));
		
		return true;

	}

	void BasicRemoteInject::Initialize()
	{
	}

	void BasicRemoteInject::Uninitialize()
	{
		PTHREAD_START_ROUTINE pfnAddr=(PTHREAD_START_ROUTINE)GetProcAddress(GetModuleHandle(TEXT("Kernel32")), "FreeLibrary");

		std::map<DWORD,SREMOTEINFO*>::iterator iter = m_mRemoteList.begin();
		while( iter != m_mRemoteList.end() )
		{
			SREMOTEINFO* pRemoteInfo = iter->second;

			if( pRemoteInfo->dwType == BasicRemoteInject::TYPE_DLL )
			{
				//HANDLE hRemoteThread =::CreateRemoteThread(pRemoteInfo->hRemoteProcess,NULL,0, (LPTHREAD_START_ROUTINE)pfnAddr,, pRemoteInfo->pVirtualAddress,0,NULL);
				HANDLE hRemoteThread =::CreateRemoteThread(pRemoteInfo->hRemoteProcess,NULL,0, (LPTHREAD_START_ROUTINE)pfnAddr, reinterpret_cast<void*>(pRemoteInfo->hLibModelAddress),0,NULL);
				::WaitForSingleObject(hRemoteThread, INFINITE);
				::CloseHandle(hRemoteThread);
			}
			//::VirtualFreeEx(pRemoteInfo->hRemoteProcess, pRemoteInfo->pVirtualAddress, pRemoteInfo->dwMemSize, MEM_RELEASE);
		}

		m_mRemoteList.clear();
	}

	bool BasicRemoteInject::RemoteInjectCode( const DWORD dwProcessId, PVOID pAddress )
	{
		DWORD THREADSIZE = 1024*2;//const DWORD THREADSIZE = 1024*4;
		HANDLE pRemoteThread = NULL;
		HANDLE hRemoteProcess = NULL;
		//PTHREAD_START_ROUTINE pfnAddr = NULL;
		void *pVirtualBase = NULL;

		hRemoteProcess=::OpenProcess(PROCESS_ALL_ACCESS,false,dwProcessId); 
		int nNum = 0;
		do 
		{
			THREADSIZE = THREADSIZE << 1;
			pVirtualBase=::VirtualAllocEx(hRemoteProcess,0,THREADSIZE,MEM_COMMIT|MEM_RESERVE,PAGE_EXECUTE_READWRITE);
		} while (++nNum < 10 && !pVirtualBase);

		if(pVirtualBase && !::WriteProcessMemory(hRemoteProcess,pVirtualBase,pAddress/*"d:\\RemoteDll.dll"*/,THREADSIZE,NULL))
		{
			::VirtualFreeEx(hRemoteProcess, pVirtualBase, THREADSIZE, MEM_RELEASE );
			return false;
		}

		DWORD tid;
		pRemoteThread=::CreateRemoteThread(hRemoteProcess,NULL,0,(LPTHREAD_START_ROUTINE)pVirtualBase,NULL,0, &tid);

		if(pRemoteThread==NULL)
		{
			::VirtualFreeEx(hRemoteProcess, pVirtualBase, THREADSIZE, MEM_RELEASE );
			return false;
		}

		SREMOTEINFO* pRemoteInfo = new SREMOTEINFO();
		pRemoteInfo->dwProcessId = dwProcessId;
		pRemoteInfo->hRemoteThread = pRemoteThread;
		pRemoteInfo->pVirtualAddress = pVirtualBase;
		pRemoteInfo->hRemoteProcess = hRemoteProcess;
		pRemoteInfo->dwMemSize = THREADSIZE;
		pRemoteInfo->dwType = BasicRemoteInject::TYPE_CODE;
		//_tcsncpy_s(pRemoteInfo->szFileName, MAX_PATH*3, szLibrary, _tcslen(szLibrary));
		this->m_mRemoteList.insert(std::make_pair(dwProcessId, pRemoteInfo));

		CloseHandle( pRemoteThread );

		return true;
	}
}
